// SPDX-FileCopyrightText: 2025 UnionTech Software Technology Co., Ltd.
// SPDX-License-Identifier: MIT
package metrics

import (
        "fmt"
        "log"
        "net"
        "io"
        "os"
        "time"
        "bufio"
        "strings"
        "strconv"
        "syscall"

        "uos-dovecot-exporter/internal/exporter"
        "github.com/prometheus/client_golang/prometheus"
)

var (
        defaultSeverVersion = "1.0.0"
        useDefaultPromReg   *bool
)

type DovecotCollector struct {
        socketPath        string
        scopes            []string
        upDesc            *prometheus.Desc
	    processCPUDesc    *prometheus.Desc
        processMemoryDesc *prometheus.Desc
        processFDDesc     *prometheus.Desc
	    maxFDs            *prometheus.Desc
        virtualMemory     *prometheus.Desc
        startTime         *prometheus.Desc
}

func init() {
        dovecotCollector := NewDovecotCollector("/var/run/dovecot/stats", []string{"user"})
        exporter.Register(dovecotCollector)
}

func NewDovecotCollector(socketPath string, scopes []string) *DovecotCollector {
        return &DovecotCollector{
                socketPath: socketPath,
                scopes:     scopes,
                upDesc: prometheus.NewDesc(
                        prometheus.BuildFQName("dovecot", "", "up"),
                        "Whether scraping Dovecot's metrics was successful.",
                        []string{"scope"},
                        nil,
                ),

                processCPUDesc: prometheus.NewDesc(
                        "custom_process_cpu_seconds_total",
                        "Total user and system CPU time spent in seconds (custom).",
                        nil, nil,
                ),

                processMemoryDesc: prometheus.NewDesc(
                        "custom_process_resident_memory_bytes",
                        "Resident memory size in bytes (custom).",
                        nil, nil,
                ),
                processFDDesc: prometheus.NewDesc(
                        "custom_process_open_fds",
                        "Number of open file descriptors (custom).",
                        nil, nil,
                ),
                maxFDs: prometheus.NewDesc(
                        "custom_process_max_fds",
                        "Maximum number of open file descriptors (custom).",
                        nil, nil,
                ),
                virtualMemory: prometheus.NewDesc(
                        "custom_process_virtual_memory_bytes",
                        "Virtual memory size in bytes (custom).",
                        nil, nil,
                ),
                startTime: prometheus.NewDesc(
                        "custom_process_start_time_seconds",
                        "Start time of the process since unix epoch in seconds (custom).",
                        nil, nil,
                ),
        }
}

func (c *DovecotCollector) Describe(ch chan<- *prometheus.Desc) {
        ch <- c.upDesc
        ch <- c.processCPUDesc
        ch <- c.processMemoryDesc
        ch <- c.processFDDesc

        ch <- c.maxFDs
        ch <- c.virtualMemory
        ch <- c.startTime
}

func (c *DovecotCollector) Collect(ch chan<- prometheus.Metric) {
        for _, scope := range c.scopes {
                log.Printf("scope: %s", scope)
                err := c.collectFromSocket(scope, ch)
                if err != nil {
                        log.Printf("Failed to scrape socket: %s", err)
                        ch <- prometheus.MustNewConstMetric(
                                c.upDesc,
                                prometheus.GaugeValue,
                                0.0,
                                scope,
                        )
                } else {
                        ch <- prometheus.MustNewConstMetric(
                                c.upDesc,
                                prometheus.GaugeValue,
                                1.0,
                                scope,
                        )
                }
        }
        if err := c.collectProcessMetrics(ch); err != nil {
                log.Printf("Process metrics collection failed: %s", err)
        }

        pid := os.Getpid()

        fds := countOpenFiles(pid)
        memStats := getProcessMemory(pid)

        startTime := getProcessStartTime(pid)

        ch <- prometheus.MustNewConstMetric(c.maxFDs, prometheus.GaugeValue, float64(fds.max))
        ch <- prometheus.MustNewConstMetric(c.virtualMemory, prometheus.GaugeValue, float64(memStats.virtual))
        ch <- prometheus.MustNewConstMetric(c.startTime, prometheus.GaugeValue, startTime)
}

func countOpenFiles(pid int) struct{ open, max int } {
        fds, _ := os.ReadDir(fmt.Sprintf("/proc/%d/fd", pid))

        var rLimit syscall.Rlimit
        if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
                log.Printf("Failed to get rlimit: %v, using default value 1024", err)
                return struct{ open, max int }{
                        open: len(fds),
                        max:  1024,
                }
        }

        return struct{ open, max int }{
                open: len(fds),
                max:  int(rLimit.Cur),
        }
}

func (c *DovecotCollector) collectProcessMetrics(ch chan<- prometheus.Metric) error {
        if stat, err := os.Open("/proc/self/stat"); err == nil {
                defer stat.Close()
                scanner := bufio.NewScanner(stat)
                scanner.Scan()
                fields := strings.Fields(scanner.Text())

                utime, _ := strconv.ParseFloat(fields[13], 64)
                stime, _ := strconv.ParseFloat(fields[14], 64)
                clockTicks := float64(100)

                ch <- prometheus.MustNewConstMetric(
                        c.processCPUDesc,
                        prometheus.CounterValue,
                        (utime+stime)/clockTicks,
                )
        }

        if status, err := os.Open("/proc/self/status"); err == nil {
                defer status.Close()
                scanner := bufio.NewScanner(status)
                for scanner.Scan() {
                        if strings.HasPrefix(scanner.Text(), "VmRSS:") {
                                kb, _ := strconv.ParseFloat(strings.Fields(scanner.Text())[1], 64)
                                ch <- prometheus.MustNewConstMetric(
                                        c.processMemoryDesc,
                                        prometheus.GaugeValue,
                                        kb*1024,
                                )
                        }
                }
        }

        if fds, err := os.ReadDir("/proc/self/fd"); err == nil {
                ch <- prometheus.MustNewConstMetric(
                        c.processFDDesc,
                        prometheus.GaugeValue,
                        float64(len(fds)),
                )
        }

	return nil
}

func (c *DovecotCollector) collectFromSocket(scope string,
        ch chan<- prometheus.Metric) error {

        var conn net.Conn
        var err error

        for i := 0; i < 3; i++ {
                conn, err = net.Dial("unix", c.socketPath)
                if err == nil {
                        break
                }
                time.Sleep(100 * time.Millisecond)
        }

        if err != nil {
                return fmt.Errorf("failed to connect to socket after 3 attempts: %w", err)
        }
        defer conn.Close()

        conn, err = net.Dial("unix", c.socketPath)
        if err != nil {
                return fmt.Errorf("failed to connect to socket: %w", err)
        }
        defer conn.Close()

        _, err = conn.Write([]byte("EXPORT\t" + scope + "\n"))
        if err != nil {
                return fmt.Errorf("failed to write to socket: %w", err)
        }

        return c.collectFromReader(conn, scope, ch)
}

func (c *DovecotCollector) collectFromReader(reader io.Reader, scope string, ch chan<- prometheus.Metric) error {
        scope string,
        ch chan<- prometheus.Metric) error {
        if scope == "global" {
                return c.collectGlobalMetrics(reader, scope, ch)
        }
        return c.collectDetailMetrics(reader, scope, ch)
}

func (c *DovecotCollector) collectDetailMetrics(reader io.Reader, _ string, ch chan<- prometheus.Metric) error {
        scanner := bufio.NewScanner(reader)
        scanner.Split(bufio.ScanLines)

        if !scanner.Scan() {
            return fmt.Errorf("failed to extract columns from input")
        }

        columnNames := strings.Split(scanner.Text(), "\t")
        if len(columnNames) < 2 {
                return fmt.Errorf("input does not provide any columns")
        }

        columns := make([]*prometheus.Desc, len(columnNames)-1)
        for i, columnName := range columnNames[1:] {
                columns[i] = prometheus.NewDesc(
                        prometheus.BuildFQName("dovecot", columnNames[0], columnName),
                        "Help text not provided by this exporter.",
                        []string{columnNames[0]},
                        nil,
                )
        }

        for scanner.Scan() {
                row := scanner.Text()
                if strings.TrimSpace(row) == "" {
                        break
                }

                values := strings.Split(row, "\t")
                if len(values) != len(columns)+1 {
                        return fmt.Errorf("value count does not match column count")
                }

                if values[0] == "" {
                        values[0] = "empty_user"
                }

                for i, value := range values[1:] {
                        f, err := strconv.ParseFloat(value, 64)
                        if err != nil {
                                return fmt.Errorf("failed to parse value: %w", err)
                        }
                        ch <- prometheus.MustNewConstMetric(
                                columns[i],
                                prometheus.UntypedValue,
                                f,
                                values[0],
                        )
                }
        }

        return scanner.Err()
}

func (c *DovecotCollector) collectGlobalMetrics(reader io.Reader,
        scope string,
        ch chan<- prometheus.Metric) error {
        scanner := bufio.NewScanner(reader)
        scanner.Split(bufio.ScanLines)

        if !scanner.Scan() {
                return fmt.Errorf("failed to extract columns from input")
        }

        columnNames := strings.Fields(scanner.Text())
        if len(columnNames) < 1 {
                return fmt.Errorf("input does not provide any columns")
        }

        columns := make([]*prometheus.Desc, len(columnNames))
        for i, columnName := range columnNames {
                columns[i] = prometheus.NewDesc(
                        prometheus.BuildFQName("dovecot", scope, columnName),
                        "Help text not provided by this exporter.",
                        []string{},
                        nil,
                )
        }

        if !scanner.Scan() {
                return scanner.Err()
        }

        values := strings.Fields(scanner.Text())
        if len(values) != len(columns) {
                return fmt.Errorf("value count does not match column count")
        }

        for i, value := range values {
                f, err := strconv.ParseFloat(value, 64)
                if err != nil {
                        return fmt.Errorf("failed to parse value: %w", err)
                }
                ch <- prometheus.MustNewConstMetric(
                        columns[i],
                        prometheus.UntypedValue,
                        f,
                )
        }

        return scanner.Err()
}

func getProcessMemory(pid int) struct{ virtual, resident uint64 } {
        data, _ := os.ReadFile(fmt.Sprintf("/proc/%d/statm", pid))
        fields := strings.Fields(string(data))

        pageSize := uint64(os.Getpagesize())
        vms, _ := strconv.ParseUint(fields[0], 10, 64)
        rss, _ := strconv.ParseUint(fields[1], 10, 64)

        return struct{ virtual, resident uint64 }{
                virtual:  vms * pageSize,
                resident: rss * pageSize,
        }
}

func getProcessStartTime(pid int) float64 {
        stat, _ := os.Stat(fmt.Sprintf("/proc/%d", pid))
        return float64(stat.ModTime().Unix())
}
